import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = nn.BatchNorm1d(4, momentum=0.2)

    def forward(self, input_x):
        out = self.op(input_x)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    labels = torch.FloatTensor(np.random.random_sample([2,4]).astype(np.float32))
    labels = torch.tensor(np.array([[0.7, 0.5, 0.5, 0.6],
                                     [0.5, 0.4, 0.6, 0.9]]).astype(np.float32))

    result = net(labels)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)


